import csv
from mmsdk import mmdatasdk

# 指定 MOSEI 的情绪标签文件
emotions_csd_path = 'CMU_MOSEI_Labels/CMU_MOSEI_Labels_Emotions.csd'

# 加载数据
emotions_dataset = mmdatasdk.mmdataset({'emotion': emotions_csd_path})

# 获取所有视频片段的情绪标签
emotion_data = emotions_dataset['emotion']['features']

# 映射字典，映射情绪到数字标签
emotion_mapping = {
    "Anger": 0,  # 愤怒
    "Disgust": 1,  # 厌恶
    "Fear": 2,  # 恐惧
    "Happy": 3,  # 幸福
    "Sad": 4,  # 悲伤
    "Surprise": 5  # 惊讶
}

# 创建一个列表保存结果
video_labels = []

# 遍历情绪数据，匹配每个视频片段的情绪标签
for video_id, emotions in emotion_data.items():
    # 将情绪标签转换为0-5的数字
    label = [emotion_mapping.get(emotion, -1) for emotion in emotions]

    # 取最大值，表示该视频片段的主导情绪
    max_emotion_label = max(label)

    # 将视频ID和对应的情绪标签添加到列表
    video_labels.append([video_id, max_emotion_label])

# 保存为CSV文件
output_file = 'video_emotion_labels.csv'
with open(output_file, mode='w', newline='', encoding='utf-8') as file:
    writer = csv.writer(file)

    # 写入CSV的头部
    writer.writerow(['video_id', 'label'])

    # 写入数据
    writer.writerows(video_labels)

print(f"视频对应情绪标签已保存为 '{output_file}'")
